"""
在 MovieLens 数据集上进行的仿真实验
"""
import random
import UCFRS


def get_all_users(path='./ml-latest-small'):
    users = set()
    is_first_line = True
    for line in open(path + '/ratings.csv'):
        # 跳过首行
        if is_first_line:
            is_first_line = False
            continue
        users.add(line.split(',')[0])
    return users


def get_all_movies(path='./ml-latest-small'):
    movies = {}
    is_first_line = True
    for line in open(path + '/movies.csv', encoding='utf-8'):
        # 跳过首行
        if is_first_line:
            is_first_line = False
            continue
        (id, title) = line.split(',')[0:2]
        movies[id] = title  # 把 title 和 id 对应
    return movies


def get_all_ratings(path='./ml-latest-small'):
    ratings = []
    is_first_line = True
    for line in open(path + '/ratings.csv'):
        # 跳过首行
        if is_first_line:
            is_first_line = False
            continue
        ratings.append(line.split(',')[0:3])
    return ratings


def split_raintgs_into_train_and_test(ratings, train_ratio, randseed):
    random.seed(randseed)
    random.shuffle(ratings)
    train_size = int(train_ratio * len(ratings))
    train = ratings[:train_size]
    test = ratings[train_size:]
    return train, test


def rating2prefs(ratings, movies):
    prefs = {}
    for rating in ratings:
        (userId, movieId, rating) = rating
        prefs.setdefault(userId, {})  # 如果还没有 userId 的评分，创建空字典；如果 userId 已经有值了，就不改变这个字典
        prefs[userId][movies[movieId]] = float(rating)
    return prefs


def experiment(seed, train_ratio, comparison_range, similarity):
    movies = get_all_movies()
    ratings = get_all_ratings()
    train_ratings, test_ratings = split_raintgs_into_train_and_test(ratings, train_ratio, seed)
    train_prefs = rating2prefs(train_ratings, movies)
    test_ratings = test_ratings[:min(len(test_ratings), comparison_range)]  # 受限于算力资源，取测试集的子集作为评测对象

    predicated_ratings = []
    mae = 0  # 预测值和测试集中的值的平均绝对误差
    cnt = 0  # 有效预测值的数量
    for rating in test_ratings:
        person = rating[0]
        movieId = rating[1]
        pred = UCFRS.predicate(train_prefs, person, movies[movieId], similarity)  # 预测值
        predicated_ratings.append((person, movieId, pred))

        if pred != 0 and float(rating[2]) != 0:
            mae += abs(pred - float(rating[2]))  # 预测值和实际值的差
            cnt += 1
    mae /= cnt

    return mae


if __name__ == '__main__':
    seed = 0
    comparison_range = 500
    train_ratios = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    similarities = [UCFRS.sim_distance, UCFRS.sim_pearson, UCFRS.sim_cosine]

    for ratio in train_ratios:
        for simi in similarities:
            mae = experiment(seed, ratio, comparison_range, simi)
            print(f"{ratio=}, {simi.__name__=}, {mae=}")
